package com.tj.collection.filters;

/**
 * @author moss
 * 2018-12-12
 */

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

public class ResponseWrapper extends HttpServletResponseWrapper {

    /**
     * This class implements an output stream in which the data is written into a byte array.
     * The buffer automatically grows as data is written to it. The data can be retrieved using toByteArray() and toString().
     * Closing a ByteArrayOutputStream has no effect. The methods in this class can be called after the stream has been closed without generating an IOException.
     */
    private ByteArrayOutputStream buffer = null;//输出到byte array
    private ServletOutputStream out = null;
    private PrintWriter writer = null;

    public ResponseWrapper(HttpServletResponse resp) throws IOException {
        super(resp);
        buffer = new ByteArrayOutputStream();// 真正存储数据的流
        out = new WapperedOutputStream(buffer);
        writer = new PrintWriter(new OutputStreamWriter(buffer, "UTF-8"));
    }

    /**
     * 重载父类获取outputstream的方法
     */
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return out;
    }

    /**
     * 重载父类获取writer的方法
     */
    @Override
    public PrintWriter getWriter() throws UnsupportedEncodingException {
        return writer;
    }

    /**
     * 重载父类获取flushBuffer的方法
     */
    @Override
    public void flushBuffer() throws IOException {
        if (out != null) {
            out.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }

    @Override
    public void reset() {
        buffer.reset();
    }

    /**
     * 将out、writer中的数据强制输出到WapperedResponse的buffer里面，否则取不到数据
     */
    public byte[] getResponseData() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }

    /**
     * 内部类，对ServletOutputStream进行包装
     */
    private class WapperedOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos = null;

        public WapperedOutputStream(ByteArrayOutputStream stream) throws IOException {
            bos = stream;
        }

        @Override
        public void write(int b) throws IOException {
            bos.write(b);
        }

        @Override
        public void write(byte[] b) throws IOException {
            bos.write(b, 0, b.length);
        }

        /**
         * Checks if a non-blocking write will succeed. If this returns
         * <code>false</code>, it will cause a callback to
         * {@link WriteListener#onWritePossible()} when the buffer has emptied. If
         * this method returns <code>false</code> no further data must be written
         * until the contain calls {@link WriteListener#onWritePossible()}.
         *
         * @return <code>true</code> if data can be written, else <code>false</code>
         * @since Servlet 3.1
         */
        @Override
        public boolean isReady() {
            return false;
        }

        /**
         * Sets the {@link WriteListener} for this {@link ServletOutputStream} and
         * thereby switches to non-blocking IO. It is only valid to switch to
         * non-blocking IO within async processing or HTTP upgrade processing.
         *
         * @param listener The non-blocking IO write listener
         * @throws IllegalStateException If this method is called if neither
         *                               async nor HTTP upgrade is in progress or
         *                               if the {@link WriteListener} has already
         *                               been set
         * @throws NullPointerException  If listener is null
         * @since Servlet 3.1
         */
        @Override
        public void setWriteListener(WriteListener listener) {

        }
    }

}